package com.huawei.aiosproj.engine;

import android.content.Context;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.text.TextUtils;
import android.util.Base64;
import android.util.Log;

import okhttp3.*;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.LinkedList;
import java.util.ListIterator;
import java.util.concurrent.TimeUnit;

public class LLMEngineRequest {
    private static final String TAG = "LLMEngineRequest";
    // Ensure this is your actual API endpoint
    private static final String ARK_API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions";
    private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");

    // Max number of *historical* user messages WITH images to include in the request.
    // The current turn's image is separate, making it up to 5 images total.
    private static final int MAX_HISTORICAL_USER_IMAGES_IN_REQUEST = 4;

    // Maximum number of image-containing user messages to RETAIN in the overall history.
    private static final int MAX_IMAGE_MESSAGES_IN_OVERALL_HISTORY = 10;


    private final Context context;
    private final OkHttpClient client;
    private final String apiKey;
    private final LinkedList<JSONObject> messageHistory; // Stores the complete conversation history

    public LLMEngineRequest(Context context, String apiKey) {
        this.context = context.getApplicationContext();
        this.apiKey = apiKey;
        this.messageHistory = new LinkedList<>();

        if (TextUtils.isEmpty(apiKey)) {
            Log.e(TAG, "API Key is null or empty. Requests will likely fail.");
        }

        this.client = new OkHttpClient.Builder()
                .connectTimeout(60, TimeUnit.SECONDS) // Increased timeout
                .readTimeout(120, TimeUnit.SECONDS)  // Increased timeout
                .writeTimeout(120, TimeUnit.SECONDS) // Increased timeout
                .build();
    }

    public String readTxtFromAssets(Context context, String fileName) {
        try (InputStream is = context.getAssets().open(fileName);
             BufferedReader reader = new BufferedReader(new InputStreamReader(is))) {
            StringBuilder sb = new StringBuilder();
            String line;
            while ((line = reader.readLine()) != null) {
                sb.append(line).append("\n");
            }
            // Remove last newline if present
            if (sb.length() > 0 && sb.charAt(sb.length() - 1) == '\n') {
                sb.deleteCharAt(sb.length() - 1);
            }
            return sb.toString();
        } catch (IOException e) {
            Log.e(TAG, "Error reading txt from assets: " + fileName, e);
            return null;
        }
    }

    public void resetConversation(String initialSystemPrompt) throws JSONException {
        messageHistory.clear();
        if (!TextUtils.isEmpty(initialSystemPrompt)) {
            JSONObject systemMessage = new JSONObject();
            // Crucially, the system prompt itself is a "user" role message with plain text content
            // according to your Python example.
            systemMessage.put("role", "user");
            systemMessage.put("content", initialSystemPrompt);
            messageHistory.add(systemMessage);
            Log.d(TAG, "Conversation reset and initial system prompt added.");
        } else {
            Log.d(TAG, "Conversation reset. No initial prompt added.");
        }
    }

    private JSONObject createTextContent(String text) throws JSONException {
        JSONObject textPart = new JSONObject();
        textPart.put("type", "text");
        textPart.put("text", text);
        return textPart;
    }

    private JSONObject createImageContent(String base64Image, String imagePath) throws JSONException {
        JSONObject imagePart = new JSONObject();
        imagePart.put("type", "image_url");
        JSONObject imageUrlData = new JSONObject();
        String imageFormat = "png"; // Default
        if (imagePath != null) {
            String lowerPath = imagePath.toLowerCase();
            if (lowerPath.endsWith(".jpg") || lowerPath.endsWith(".jpeg")) {
                imageFormat = "jpeg";
            } else if (lowerPath.endsWith(".png")) {
                imageFormat = "png";
            } else if (lowerPath.endsWith(".webp")) {
                imageFormat = "webp";
            }
            // Add more formats if needed
        }
        imageUrlData.put("url", "data:image/" + imageFormat + ";base64," + base64Image);
        imagePart.put("image_url", imageUrlData);
        return imagePart;
    }

    private void addMessageToInternalHistory(String role, Object content) throws JSONException {
        JSONObject message = new JSONObject();
        message.put("role", role);
        message.put("content", content); // content can be String or JSONArray
        this.messageHistory.add(message);
    }

    private boolean isImageMessage(JSONObject message) {
        if (!"user".equals(message.optString("role"))) return false;
        Object content = message.opt("content");
        if (content instanceof JSONArray) {
            JSONArray contentArray = (JSONArray) content;
            for (int i = 0; i < contentArray.length(); i++) {
                JSONObject part = contentArray.optJSONObject(i);
                if (part != null && "image_url".equals(part.optString("type"))) {
                    return true;
                }
            }
        }
        return false;
    }

    private String imageToBase64(String imagePath) throws IOException {
        if (TextUtils.isEmpty(imagePath)) {
            Log.w(TAG, "imageToBase64: imagePath is null or empty.");
            return null;
        }
        File imageFile = new File(imagePath);
        if (!imageFile.exists() || !imageFile.canRead()) {
            Log.e(TAG, "imageToBase64: Image file does not exist or cannot be read: " + imagePath);
            return null;
        }

        Bitmap bitmap = BitmapFactory.decodeFile(imagePath);
        if (bitmap == null) {
            Log.e(TAG, "imageToBase64: Failed to decode image file: " + imagePath);
            return null;
        }

        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        String extension = "";
        int lastDot = imagePath.lastIndexOf('.');
        if (lastDot > 0 && lastDot < imagePath.length() - 1) {
            extension = imagePath.substring(lastDot + 1).toLowerCase();
        }

        Bitmap.CompressFormat format = Bitmap.CompressFormat.PNG; // Default
        if ("jpg".equals(extension) || "jpeg".equals(extension)) {
            format = Bitmap.CompressFormat.JPEG;
        } else if ("webp".equals(extension)) {
            format = Bitmap.CompressFormat.WEBP;
        }
        // PNG is already default

        bitmap.compress(format, 100, byteArrayOutputStream); // Quality 100
        byte[] byteArray = byteArrayOutputStream.toByteArray();
        bitmap.recycle(); // Important to free memory
        return Base64.encodeToString(byteArray, Base64.NO_WRAP);
    }


    /**
     * Constructs the "messages" array for the API request based on the UI-TARS "history-5" logic:
     * 1. System Prompt (user role)
     * 2. All historical assistant responses.
     * 3. The most recent MAX_HISTORICAL_USER_IMAGES_IN_REQUEST (e.g., 4) historical user images.
     *    (Each user image is followed by its corresponding assistant response if available in history).
     * 4. The current user's image/prompt.
     *
     * @param currentImagePath  Path to the image for the current user turn.
     * @param currentUserPrompt Text prompt for the current user turn.
     * @return JSONArray of messages formatted for the API.
     * @throws JSONException If JSON construction fails.
     * @throws IOException   If image encoding fails.
     */
    private JSONArray buildMessagesForRequest(String currentImagePath, String currentUserPrompt) throws JSONException, IOException {
        JSONArray messagesToSend = new JSONArray();

        // 0. Construct the current user's message (not yet in main history)
        JSONObject currentUserTurnMessage = new JSONObject();
        currentUserTurnMessage.put("role", "user");
        JSONArray currentUserContentArray = new JSONArray();
        boolean currentTurnHasContent = false;
        if (!TextUtils.isEmpty(currentImagePath)) {
            String base64Image = imageToBase64(currentImagePath);
            if (base64Image == null) {
                throw new IOException("Failed to encode current image to Base64: " + currentImagePath);
            }
            currentUserContentArray.put(createImageContent(base64Image, currentImagePath));
            currentTurnHasContent = true;
        }
        if (!TextUtils.isEmpty(currentUserPrompt)) {
            currentUserContentArray.put(createTextContent(currentUserPrompt));
            currentTurnHasContent = true;
        }
        if (!currentTurnHasContent) { // API might require at least one content part for user
            Log.w(TAG, "Current user turn has no image and no prompt. Sending empty content for user.");
            // Depending on API: send empty array or a default minimal text part.
            // For example: currentUserContentArray.put(createTextContent(""));
            currentUserTurnMessage.put("content", new JSONArray()); // As per your Python example for image-only turns
        } else {
            currentUserTurnMessage.put("content", currentUserContentArray);
        }

        // --- Start building messagesToSend ---

        // 1. Add System Prompt (always the first message in messageHistory)
        if (messageHistory.isEmpty()) {
            Log.e(TAG, "Message history is empty, cannot build request. Reset conversation first.");
            // Potentially throw an error or handle this state appropriately
            if (currentUserTurnMessage.has("content")) { // If only current message, send it.
                messagesToSend.put(currentUserTurnMessage);
                return messagesToSend;
            }
            return new JSONArray(); // Or throw
        }
        JSONObject systemPrompt = messageHistory.getFirst();
        if ("user".equals(systemPrompt.optString("role")) && systemPrompt.opt("content") instanceof String) {
            messagesToSend.put(systemPrompt);
        } else {
            Log.w(TAG, "First message in history is not a valid system prompt. Skipping.");
            // This case should ideally not happen if resetConversation is used correctly.
        }

        // 2. Identify historical user images and all assistant messages
        LinkedList<JSONObject> historicalUserImageMessages = new LinkedList<>();
        LinkedList<JSONObject> allAssistantMessages = new LinkedList<>();

        // Iterate through history (skipping system prompt, which is already added)
        for (int i = 1; i < messageHistory.size(); i++) {
            JSONObject msg = messageHistory.get(i);
            String role = msg.optString("role");

            if ("user".equals(role) && isImageMessage(msg)) {
                historicalUserImageMessages.add(msg);
            } else if ("assistant".equals(role)) {
                allAssistantMessages.add(msg);
            }
        }

        // 3. Add all historical assistant messages that appear *before* the selected image block
        //    AND the assistant messages that are part of the selected image block.
        //    The logic is: system_prompt, then all relevant assistant messages, then user_image_1, then its_assistant_1 (if exists),
        //    user_image_2, its_assistant_2, ..., current_user_image
        //    A simpler interpretation from your example: Sys, Asst, Asst, Asst, UserImg, Asst, UserImg, Asst ...

        // According to your python example, it's:
        // System Prompt
        // Assistant (round 1)
        // Assistant (round 2)
        // Assistant (round 3)
        // User (round 4 image)
        // Assistant (round 4 response)
        // ...
        // User (round N image - current turn)

        // This means we add assistant messages based on their original position relative to user images.

        LinkedList<JSONObject> tempFullHistory = new LinkedList<>(messageHistory); // Use a copy of full history
        // up to before current turn for ordering
        int userImagesForRequestCount = 0;
        int lastAddedHistoryIndex = 0; // Index in tempFullHistory

        // Add assistant messages that are *not* directly part of an image turn that will be skipped
        // More accurately, we collect all assistant messages and the *last N* image turns.

        // Get the last MAX_HISTORICAL_USER_IMAGES_IN_REQUEST from historicalUserImageMessages
        LinkedList<JSONObject> recentHistoricalUserImages = new LinkedList<>();
        ListIterator<JSONObject> iter = historicalUserImageMessages.listIterator(historicalUserImageMessages.size());
        while(iter.hasPrevious() && recentHistoricalUserImages.size() < MAX_HISTORICAL_USER_IMAGES_IN_REQUEST) {
            recentHistoricalUserImages.addFirst(iter.previous());
        }

        // Now, iterate through the original messageHistory (after system prompt)
        // and add messages according to the desired structure.
        for (int i = 1; i < messageHistory.size(); i++) { // Start after system prompt
            JSONObject historyMsg = messageHistory.get(i);
            String role = historyMsg.optString("role");

            if ("assistant".equals(role)) {
                messagesToSend.put(historyMsg);
            } else if ("user".equals(role) && isImageMessage(historyMsg)) {
                // Only add this user image if it's in our `recentHistoricalUserImages` list
                if (recentHistoricalUserImages.contains(historyMsg)) {
                    messagesToSend.put(historyMsg);
                    // Check for its corresponding assistant response immediately after it in original history
                    if (i + 1 < messageHistory.size()) {
                        JSONObject nextMsg = messageHistory.get(i + 1);
                        if ("assistant".equals(nextMsg.optString("role"))) {
                            // This assistant message is already added by the general "assistant" role check above.
                            // The key is that the main loop for "assistant" adds ALL of them.
                            // The check for "user" image messages from `recentHistoricalUserImages` ensures only those are added.
                            // The ordering is implicitly handled by iterating through `messageHistory`.
                        }
                    }
                } else {
                    // This is an older user image message, not part of the recent N. Skip it.
                    // Its assistant response (if any) would have been added by the assistant role check,
                    // which matches your Python example (all assistant responses are included).
                }
            }
            // Non-image user messages (if any) are skipped as per your history-5 logic.
        }

        // 4. Add the current user's message (image and/or prompt)
        if (currentUserTurnMessage.has("content")) { // Ensure it was properly formed
            messagesToSend.put(currentUserTurnMessage);
        }

        Log.i(TAG, "Prepared messages for API. Count: " + messagesToSend.length());
        if (Log.isLoggable(TAG, Log.VERBOSE)) { // Use VERBOSE for potentially large output
            try {
                Log.v(TAG, "Messages to send JSON: " + messagesToSend.toString(2));
            } catch (JSONException e) { /* ignore */ }
        }
        return messagesToSend;
    }


    static int requestCount = 0; // For unique log file names

    public String runModel(String imagePath, String userPrompt) throws IOException, JSONException {
        // 1. Construct the messages to send based on history and current input
        JSONArray messagesForApi;
        JSONObject currentUserMessageForHistory = new JSONObject(); // This will be added to main history
        currentUserMessageForHistory.put("role", "user");
        JSONArray currentUserContentArray = new JSONArray();

        boolean hasContentForCurrentTurn = false;
        if (!TextUtils.isEmpty(imagePath)) {
            String base64Image = imageToBase64(imagePath);
            if (base64Image == null) {
                throw new IOException("Failed to encode image to Base64 from path: " + imagePath);
            }
            currentUserContentArray.put(createImageContent(base64Image, imagePath));
            hasContentForCurrentTurn = true;
        }
        if (!TextUtils.isEmpty(userPrompt)) {
            currentUserContentArray.put(createTextContent(userPrompt));
            hasContentForCurrentTurn = true;
        }

        if (!hasContentForCurrentTurn) {
            Log.w(TAG, "runModel called with no image and no user prompt for current turn.");
            // API might require content. If so, create an empty text part or similar.
            // As per python, image-only turns have a content array with just the image.
            // If both image and prompt are null, send empty content array for the user role.
            currentUserMessageForHistory.put("content", new JSONArray());
        } else {
            currentUserMessageForHistory.put("content", currentUserContentArray);
        }


        // Build the messages array for the API call.
        messagesForApi = buildMessagesForRequest(imagePath, userPrompt);


        // 2. Build request body
        JSONObject requestBody = new JSONObject();
        // Replace with your actual model name if different
        requestBody.put("model", "doubao-1-5-ui-tars-250428");
        requestBody.put("temperature", 0.0); // As per typical UI-TARS settings
        requestBody.put("messages", messagesForApi);

        Log.i(TAG, "Sending request to LLM. Number of messages in request: " + messagesForApi.length());
        if (Log.isLoggable(TAG, Log.VERBOSE)) {
            try {
                Log.v(TAG, "Request JSON: " + requestBody.toString(2));
            } catch (JSONException e) {
                Log.v(TAG, "Request JSON (raw): " + requestBody.toString());
            }
        }

        // Save request for debugging
        if (context != null && context.getFilesDir() != null) {
            File requestFile = new File(context.getFilesDir(), "request_" + System.currentTimeMillis() + "_" + requestCount + ".json");
            try (FileOutputStream fos = new FileOutputStream(requestFile);
                 OutputStreamWriter osw = new OutputStreamWriter(fos, "UTF-8")) {
                osw.write(requestBody.toString(2)); // pretty print
                osw.flush();
                Log.d(TAG, "Saved request to: " + requestFile.getAbsolutePath());
            } catch (IOException | JSONException e) {
                Log.w(TAG, "Failed to save request JSON", e);
            }
            requestCount++;
        } else {
            Log.w(TAG, "Context or FilesDir is null, cannot save request JSON.");
        }


        // 3. Create HTTP request
        Request request = new Request.Builder()
                .url(ARK_API_URL)
                .addHeader("Authorization", "Bearer " + this.apiKey)
                .post(RequestBody.create(requestBody.toString(), JSON))
                .build();

        // 4. Execute and process response
        try (Response response = client.newCall(request).execute()) {
            Log.e(TAG, "runModel original response:  " + response.toString() );
            String responseBodyString = response.body() != null ? response.body().string() : null;

            if (!response.isSuccessful()) {
                Log.e(TAG, "API Request Failed. Code: " + response.code() + " Body: " + responseBodyString);
                throw new IOException("Unexpected API response code: " + response.code() + " with body: " + responseBodyString);
            }

            if (responseBodyString == null) {
                Log.e(TAG, "API Response body is null on success");
                throw new IOException("API returned successful code but with empty response body.");
            }

            Log.d(TAG, "API Response Success (first 200 chars): " + responseBodyString.substring(0, Math.min(responseBodyString.length(), 200)));

            // If successful, NOW add the current user's message to the *main* messageHistory
            // Ensure `currentUserMessageForHistory` accurately reflects what was sent
            addMessageToInternalHistory("user", currentUserMessageForHistory.get("content"));
            Log.d(TAG, "Added current user message to main history. Size: " + messageHistory.size());


            // 5. Parse assistant's response and add it to main messageHistory
            try {
                JSONObject fullApiResponse = new JSONObject(responseBodyString);
                if (fullApiResponse.has("choices")) {
                    JSONArray choices = fullApiResponse.getJSONArray("choices");
                    if (choices.length() > 0) {
                        JSONObject firstChoice = choices.getJSONObject(0);
                        if (firstChoice.has("message")) {
                            JSONObject assistantMessageObj = firstChoice.getJSONObject("message");
                            if (assistantMessageObj.has("role") && assistantMessageObj.has("content")) {
                                String assistantRole = assistantMessageObj.getString("role");
                                Object assistantContent = assistantMessageObj.get("content"); // Usually String for assistant
                                addMessageToInternalHistory(assistantRole, assistantContent);
                                Log.d(TAG, "Added assistant response to main history. New size: " + messageHistory.size());
                            } else {
                                Log.w(TAG, "Assistant message in API response missing 'role' or 'content'.");
                            }
                        } else {
                            Log.w(TAG, "First choice in API response missing 'message' field.");
                        }
                    } else {
                        Log.w(TAG, "'choices' array in API response is empty.");
                    }
                } else if (fullApiResponse.has("error")) {
                    Log.e(TAG, "API returned an error in response body: " + fullApiResponse.getJSONObject("error").toString());
                } else {
                    Log.w(TAG, "Could not find 'choices' or 'error' in API response.");
                }
            } catch (JSONException e) {
                Log.e(TAG, "Failed to parse API response JSON or add assistant message to history", e);
            }

            pruneOverallHistory();
            Log.d(TAG, "Overall history pruned. Size after pruning: " + messageHistory.size());

            return responseBodyString;

        }
    }


    /**
     * Prunes the overall messageHistory to prevent it from growing indefinitely.
     * Keeps the system prompt, all assistant messages, and a certain number of
     * recent image-containing user messages.
     */
    private void pruneOverallHistory() {
        if (messageHistory.isEmpty()) {
            return;
        }

        LinkedList<JSONObject> userImageMessagesInHistory = new LinkedList<>();
        for (JSONObject msg : messageHistory) {
            if (isImageMessage(msg)) {
                userImageMessagesInHistory.add(msg);
            }
        }

        if (userImageMessagesInHistory.size() <= MAX_IMAGE_MESSAGES_IN_OVERALL_HISTORY) {
            return; // No need to prune image messages
        }

        int imagesToRemoveCount = userImageMessagesInHistory.size() - MAX_IMAGE_MESSAGES_IN_OVERALL_HISTORY;
        LinkedList<JSONObject> imagesToRemove = new LinkedList<>();
        for (int i = 0; i < imagesToRemoveCount; i++) {
            imagesToRemove.add(userImageMessagesInHistory.get(i)); // Oldest ones
        }

        LinkedList<JSONObject> newHistory = new LinkedList<>();
        JSONObject systemPrompt = null;
        if (!messageHistory.isEmpty() && "user".equals(messageHistory.getFirst().optString("role")) && messageHistory.getFirst().opt("content") instanceof String) {
            systemPrompt = messageHistory.getFirst();
            newHistory.add(systemPrompt);
        }

        for (int i = (systemPrompt != null ? 1 : 0); i < messageHistory.size(); i++) {
            JSONObject msg = messageHistory.get(i);
            if ("assistant".equals(msg.optString("role"))) {
                newHistory.add(msg); // Always keep assistant messages
            } else if (isImageMessage(msg)) {
                if (!imagesToRemove.contains(msg)) {
                    newHistory.add(msg); // Keep if it's not in the removal list
                }
            } else {
                // Other user messages (e.g. text-only, not system prompt)
                // Decide if these should be kept or pruned based on specific rules.
                // For UI-TARS like logic, often these are not present or are part of system prompt.
                // If they can exist, add them: newHistory.add(msg);
            }
        }
        messageHistory.clear();
        messageHistory.addAll(newHistory);
        Log.d(TAG, "Pruned overall history. New size: " + messageHistory.size());
    }
}